from skimage.measure import compare_ssim
import numpy as np
import urllib.request
import cv2
import base64
import hashlib
"""
ssim和md5图片相似度计算方法
"""

# base64转img
def base64_to_img(base64_str):
    # 传入为RGB格式下的base64，传出为RGB格式的numpy矩阵
    byte_data = base64.b64decode(base64_str)  #将base64转换为二进制
    encode_image = np.asarray(bytearray(byte_data), dtype="uint8")  # 二进制转换为一维数组
    img = cv2.imdecode(encode_image, cv2.IMREAD_COLOR)  # 用cv2解码为三通道矩阵

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # BGR2RGB

    return img

def ssim(img1_input, img2_input, threshold, img_type):
    result = {}
    try:
        if img_type == "url":

            res1 = urllib.request.urlopen(img1_input)
            image1 = np.asarray(bytearray(res1.read()), dtype="uint8")
            image1 = cv2.imdecode(image1, cv2.IMREAD_COLOR)

            res2 = urllib.request.urlopen(img2_input)
            image2 = np.asarray(bytearray(res2.read()), dtype="uint8")
            image2 = cv2.imdecode(image2, cv2.IMREAD_COLOR)
            image2 = np.resize(image2, (image1.shape[0], image1.shape[1], image1.shape[2]))
        if img_type == "base64":
            image1 = base64_to_img(img1_input)

            image2 = base64_to_img(img2_input)
        #
        ssim = compare_ssim(image1, image2, multichannel=True)
        # print(ssim)
        result["msg"] = "检测成功"
        result["code"] = 0
        result["algo"] = "ssim"
        result["similar"] = ssim
        # print(ssim)
        if ssim >= threshold:
            result["res"] = "相似"
        else:
            result["res"] = "不相似"

        return result
    except:
        result["msg"] = "检测失败"
        result["code"] = -1
        return result

def md5(img1_input, img2_input, img_type):

    global image11, image21
    result = {}
    try:
        if img_type == "url":
            res1 = urllib.request.urlopen(img1_input)
            image11 = np.asarray(bytearray(res1.read()), dtype="uint8")
            image11 = cv2.imdecode(image11, cv2.IMREAD_COLOR)

            res2 = urllib.request.urlopen(img2_input)
            image21 = np.asarray(bytearray(res2.read()), dtype="uint8")
            image21 = cv2.imdecode(image21, cv2.IMREAD_COLOR)

        if img_type == "base64":
            image11 = base64_to_img(img1_input)

            image21 = base64_to_img(img2_input)
        #
        file1 = image11
        file2 = image21
        md = hashlib.md5()

        md.update(file1)
        res1 = md.hexdigest()

        md = hashlib.md5()
        md.update(file2)
        res2 = md.hexdigest()

        result["msg"] = "检测成功"
        result["code"] = 0
        result["algo"] = "md5"
        if res1 == res2:
            result["res"] = "相似"
        else:
            result["res"] = "不相似"

        return result
    except:
        result["msg"] = "检测失败"
        result["code"] = -1
        return result
